import numpy as np
import platgo as pg


def roulette_wheel_selection(N: int, pop: pg.Population = None, fitness: np.ndarray = None) -> np.ndarray:
    """
    roulette_wheel_selection(N: int, pop: pg.Population = None, fitness: np.ndarray = None) returns
    the indices of N solutions by roulette-wheel selection based on fitness. A SMALLER fitness value
    indicates a LARGER probability to be selected.
    :param N:需要选择的个体数目
    :param pop:种群
    :param fitness:多目标函数的适应度值
    :return:被选择个体的索引值
    """
    assert pop or np.all(fitness), "pop and fitness can't be both None"
    if pop is None:
        # 多目标函数应有自己的适应值
        fitness = fitness
    else:
        # 单目标的适应值就是函数目标值
        fitness = pop.objv
    # 将适应值转化为一维
    fitness = fitness.reshape(1, -1)[0]
    fitness = fitness + min(min(fitness), 0)
    fitness = 1/fitness
    fitness = fitness/np.sum(fitness)
    index = np.random.choice(np.arange(len(fitness)), N, replace=True, p=fitness)
    return index
